package io;

import java.io.BufferedReader;
import java.io.FileReader;
import java.io.IOException;
import java.util.BitSet;
import java.util.HashMap;
import java.util.Map;
import java.util.Random;
import java.util.StringTokenizer;
import java.util.Vector;
import java.util.regex.Matcher;
import java.util.regex.Pattern;

import preprocessing.WindowingSystem;

import datastructures.Individual;
import datastructures.IndividualInfo;
import datastructures.PhenotypeInfo;
import datastructures.SNPInfo;
import datastructures.UserParameters;

public class PlinkFilereader {

	public static final int MAX_INDIVIDUALS = 5000;
	public static final int MAX_GENOTYPES = 2000000;
	
		
	public Individual[] readRawFile(String rawFile) throws IOException 
	{	
		System.out.println("Reading RAW file: "+  rawFile);
		BufferedReader reader = new BufferedReader(new FileReader(rawFile));
		reader.readLine(); //waste header

		Vector<Individual> rawFileContents = new Vector<Individual>();			

		String text = null;
		StringTokenizer rawTokenizer; 
		while( (text=reader.readLine()) != null ) {			
			
			rawTokenizer = new StringTokenizer (text);
			
			String FAM_ID = new String(rawTokenizer.nextToken());
			String IND_ID = new String(rawTokenizer.nextToken());
			String PAT_ID = new String(rawTokenizer.nextToken());
			String MAT_ID = new String(rawTokenizer.nextToken());
			String Sex = new String(rawTokenizer.nextToken());
			String Phenotype = new String(rawTokenizer.nextToken());
			
			int genotype_count=0;
			BitSet genotypeBitSet = new BitSet(); 
			
			while (rawTokenizer.hasMoreTokens()) {
				
				int position = 2*genotype_count;
				
				String nextToken = rawTokenizer.nextToken();
				if ( nextToken.compareTo("0") == 0 ) 
					genotypeBitSet.clear(position, position+2);
				else if ( nextToken.compareTo("1") == 0 ) {
					genotypeBitSet.set(position, position+1); //to position is exclusive : set(from, to)
					genotypeBitSet.clear(position+1,position+2);
				}
				else if ( nextToken.compareTo("2") == 0 ) 
					genotypeBitSet.set(position, position+2); //to position is exclusive : set(from, to)
				else if ( nextToken.compareTo("NA") == 0 ) 
					genotypeBitSet.clear(position, position+2);
				
				genotype_count++;
				
			}	
			rawTokenizer = null;
			rawFileContents.add ( new Individual(FAM_ID, IND_ID, PAT_ID, MAT_ID, Sex, Phenotype, genotype_count, genotypeBitSet) );
			
		}
		
		Individual[] individuals = new Individual[rawFileContents.size()];
		rawFileContents.toArray(individuals);
		
		reader.close();
		return individuals;
		
	}
	
	
	/*
	 * Helper function : takes a raw file array and divides into two based on phenotype status.
	 * Checks if the user has also input a different phenotype file, and resolves membership accordingly.
	 * Returned: First array is cases, second array is controls.
	 * Ambiguous phenotype assigned to controls.
	 */
	public Vector<Individual[]> segregateCasesAndControls(Individual[] allIndividuals, String phenoFile, boolean permute) throws IOException, PlinkFileException
	{
		Vector<Individual> cases = new Vector<Individual>();
		Vector<Individual> controls = new Vector<Individual>();
		
		if ( !permute )
			allIndividuals = resolvePhenotypeStatus(allIndividuals, phenoFile);
		else
			allIndividuals = randomizePhenotypeStatus(allIndividuals);
		
		for(int indiv=0; indiv<allIndividuals.length; indiv++) {
			if ( allIndividuals[indiv].getIndividualInfo().getPhenotype().compareTo("2") == 0 )
				cases.add( allIndividuals[indiv] );
			else if ( allIndividuals[indiv].getIndividualInfo().getPhenotype().compareTo("1") == 0 )
				controls.add( allIndividuals[indiv] );
		}
		
		Individual[] casesArray = new Individual[cases.size()]; 
		cases.toArray(casesArray);
		
		Individual[] controlsArray = new Individual[controls.size()];
		controls.toArray(controlsArray);
		
		Vector<Individual[]> casesAndControls = new Vector<Individual[]>();
		casesAndControls.add(casesArray);
		casesAndControls.add(controlsArray);
		
		return casesAndControls;
	}
	
	public Individual[] randomizePhenotypeStatus ( Individual[] allIndividuals ) throws PlinkFileException
	{		
		int number_of_cases = 0;
		int number_of_controls = 0;
		for( Individual individual : allIndividuals) {
			if (individual.getIndividualInfo().getPhenotype().compareTo("2") == 0 )
				number_of_cases += 1;
			else if ( individual.getIndividualInfo().getPhenotype().compareTo("1") == 0 )
				number_of_controls += 1;
		}
		
		Random rand = new Random();
		double cutoff = (double) number_of_cases / (number_of_cases + number_of_controls);
		
		for( Individual individual : allIndividuals) {
			if ( rand.nextDouble() < cutoff)
				individual.getIndividualInfo().setPhenotype("2");
			else
				individual.getIndividualInfo().setPhenotype("1");
		}
		return allIndividuals;
	}
	
	private Individual[] resolvePhenotypeStatus ( Individual[] allIndividuals, String phenoFile ) throws PlinkFileException
	{
		PhenotypeInfo[] phenotypeMembership = null;
		
		if ( phenoFile.compareTo("") == 0) 
			return allIndividuals;
		else 
		{
			try {
				phenotypeMembership = readPhenofile(phenoFile);
			} catch (IOException e) {
				System.err.println("Trouble reading phenotype file ... continuing with phenotypes provided in raw file.");
				return allIndividuals;
			}
		}
		
		if ( allIndividuals.length != phenotypeMembership.length )
			throw new PlinkFileException("Phenotype file and Raw file individual IDs don't match");
		
		for(int count=0; count<allIndividuals.length; count++) 
		{
			if ( allIndividuals[count].getIndividualInfo().getFAM_ID().compareTo( phenotypeMembership[count].getFAM_ID() ) == 0
					&& allIndividuals[count].getIndividualInfo().getIND_ID().compareTo( phenotypeMembership[count].getIND_ID() ) == 0 ) 
			{
					IndividualInfo currentInfo = allIndividuals[count].getIndividualInfo();
					currentInfo.setPhenotype ( phenotypeMembership[count].getPhenotype() );
					allIndividuals[count].setIndividualInfo(currentInfo);
			}
					
		}
		
		return allIndividuals;
	}
	
	/*
	 * Helper Function - let's avoid the memory wastage.
	 */
	public static BitSet StringArrayToBitSet(String[] genotypeList) throws PlinkFileException {
		
		int numberOfSNPs;
		if(genotypeList == null)
			numberOfSNPs = 0;
		else
			numberOfSNPs = genotypeList.length;
		
		BitSet genotypeVector = new BitSet (numberOfSNPs * 2); // 2 variables per SNP, representing dominance and recessive traits.
		genotypeVector.clear();
		
		
		StringTokenizer genotypeTokenizer;
		for(int index=0; index<numberOfSNPs; index++)
		{	
			genotypeTokenizer = new StringTokenizer(genotypeList[index], " ");
			String[] alleles;
			
			String first = genotypeTokenizer.nextToken();
			if(genotypeTokenizer.hasMoreTokens()) {
				alleles = new String[] {first, genotypeTokenizer.nextToken()};
			}
			else
				alleles = new String[] {first};
			
			int position = 2*index;
			int sum;
			if( alleles.length == 1) {
				//raw file
				sum = Integer.parseInt(alleles[0]);
			}
			else if ( alleles.length == 2 ) {
				//ped file
				sum = Integer.parseInt(alleles[0]) + Integer.parseInt(alleles[1]) - 2; //11 = maj hom, 12/21 = het, 22 = min hom. 
			}	
			else
				throw new PlinkFileException("Bad file format");
			
			
			switch( sum ) {
				case 0:
					//do nothing. default bit set to 0.
					break;
					
				case 1:
					genotypeVector.set(position, position+1); //to position is exclusive : set(from, to)
					break;
					
				default:		
					genotypeVector.set(position, position+2); //to position is exclusive : set(from, to)
					break;
			}
		
		}		
		
		return genotypeVector;
	}

	
	/**
	 * @param args
	 */	
	public SNPInfo[] readMapfile(String mapFile) throws IOException 
	{
		SNPInfo[] mostRecentlyReadMapFile = null;
		Vector<SNPInfo> filecontents = new Vector<SNPInfo> ();

		System.out.println("Reading MAP file: " + mapFile);
		BufferedReader reader = new BufferedReader(new FileReader(mapFile));

		String text = null;
		while ((text = reader.readLine()) != null) {
			String[] mapLine = text.split("\\s+");
			SNPInfo snp = new SNPInfo(Integer.parseInt(mapLine[0]), mapLine[1], Double.parseDouble(mapLine[2]), Integer.parseInt(mapLine[3]));
			filecontents.add(snp);
		}
		
		//filecontents.printItemList();
		mostRecentlyReadMapFile = new SNPInfo[ filecontents.size() ];
		filecontents.toArray( mostRecentlyReadMapFile );
		
		reader.close();
		return mostRecentlyReadMapFile;
	}


	/**
	 * @param args
	 */	
	public Map<String, Double> readGeneticMapfile(String cmFile) throws IOException 
	{
		
		HashMap<String, Double> filecontents = new HashMap<String, Double>();

		System.out.println("Reading Genetic Map file: " + cmFile);
		BufferedReader reader = new BufferedReader(new FileReader(cmFile));

		//entry format : chr \t snp \t cM \t Bp
		String text = null;
		while ((text = reader.readLine()) != null) {
			String[] mapLine = text.split("\\s+");
			filecontents.put(mapLine[1], Double.parseDouble(mapLine[2])); 
		}
		
		return filecontents;
	}

	
	
	public PhenotypeInfo[] readPhenofile(String phenoFile) throws IOException 
	{
		PhenotypeInfo[] mostRecentlyReadPhenoFile = null;
		Vector<PhenotypeInfo> filecontents = new Vector<PhenotypeInfo> ();
		
		System.out.println("Reading Phenotype file: " + phenoFile);
		BufferedReader reader = new BufferedReader(new FileReader(phenoFile));

		String text = null;
		while ((text = reader.readLine()) != null) {
			String[] phenoLine = text.split("\\s+");
			PhenotypeInfo pheno = new PhenotypeInfo( phenoLine[0], phenoLine[1], phenoLine[2]);
			filecontents.add(pheno);
		}
		
		//filecontents.printItemList();
		mostRecentlyReadPhenoFile = new PhenotypeInfo[ filecontents.size() ];
		filecontents.toArray( mostRecentlyReadPhenoFile );
		
		reader.close();
		return mostRecentlyReadPhenoFile;
	}
	

	
	// List of SNPs to avoid (typically due to marginal effects)
	public String[] readAvoidFile(String snpsToAvoidFile) throws IOException {
		
		Vector<String> filecontents = new Vector<String> ();

		System.out.println("Reading SNPs to avoid file: " + snpsToAvoidFile + "...");
		BufferedReader reader = new BufferedReader(new FileReader(snpsToAvoidFile));

		String text = null;
		while ((text = reader.readLine()) != null) {
			String[] avoidSnpLine = text.split("\\s+");
			filecontents.add(avoidSnpLine[0]);
		}
		
		//filecontents.printItemList();
		String[] returnArray = new String[ filecontents.size() ];
		filecontents.toArray( returnArray );
		
		reader.close();
		return returnArray;
	}
	
	
	public int[][] readSignificanceFile(String significanceFile) throws IOException {
		Vector<int[]> filecontents = new Vector<int[]>();
		
		System.out.println("Reading Significance file: " + significanceFile + "...");
		BufferedReader reader = new BufferedReader(new FileReader (significanceFile));

		String text = null;
		while ( (text = reader.readLine()) != null ) 
		{
			String[] required = text.split("\\s+");
			
			int[] entries = new int[required.length];
			for(int i=0; i<required.length; i++)
				entries[i] = Integer.parseInt( required[i] );
			
			filecontents.add(entries);
		}
		
		int[][] returnArray = new int [filecontents.size()][];
		for(int entry=0; entry<filecontents.size(); entry++) 
			returnArray[entry] = filecontents.elementAt(entry);
			
		return returnArray;
	}
	

	public BitSet[] readFrqFile (UserParameters userparams, WindowingSystem windowMaker, SNPInfo[] snpList) throws IOException, PlinkFileException
	{		
		String frqFile = userparams.getFrqfilename();
		int freqWindows = userparams.getFreqWindows(); 

		BitSet[] masks = new BitSet[ freqWindows ];
		for ( int i=0; i < freqWindows; i++ )
			masks[i] = new BitSet(snpList.length * 2); // all zeros

		//if frequency file provided
		System.out.println("\nReading FRQ file: " + frqFile + "...");		
		BufferedReader reader = new BufferedReader(new FileReader(frqFile));

		String text = null;
		reader.readLine(); //waste the raw file header.

		int count = 0; 
		while ( (text = reader.readLine()) != null && count < snpList.length ) 
		{

			Pattern pattern = Pattern.compile("\\s+(.*?)\\s+(.*?)\\s+(.*?)\\s+(.*?)\\s+(.*?)\\s+(.*?)$");
			Matcher matcher = pattern.matcher(text);

			if (matcher.matches()) 
			{
				float maf = Float.parseFloat ( matcher.group(5) );
				
				int caseDomCarriers =  (int) (userparams.getNumber_of_cases() * (2*maf*(1-maf) + maf*maf) );		
				int caseRecCarriers = (int) (userparams.getNumber_of_cases() * maf*maf );
				snpList[count].setCaseDomCarriers(caseDomCarriers);
				snpList[count].setCaseRecCarriers(caseRecCarriers);
				
				int controlDomCarriers = (int) (userparams.getNumber_of_controls() * (2*maf*(1-maf) + maf*maf) );
				int controlRecCarriers = (int) (userparams.getNumber_of_controls() * maf*maf );
				snpList[count].setControlDomCarriers(controlDomCarriers);
				snpList[count].setControlRecCarriers(controlRecCarriers);

				int[] domThenRecIndex = windowMaker.whichMask(maf);
				snpList[count].setWhichMaskForDom(domThenRecIndex[0]);
				snpList[count].setWhichMaskForRec(domThenRecIndex[1]);
				
				if(domThenRecIndex[0] != -1)
					masks[domThenRecIndex[0]].set(2*count, 2*count+1);
				if(domThenRecIndex[1] != -1)
					masks[domThenRecIndex[1]].set(2*count+1, 2*count+2);
			} 
			else 
				throw new IOException("Could not parse ped file " + frqFile);

			count++;
		} 

		if ( reader.readLine() != null || count != snpList.length )
			throw new PlinkFileException("FRQ and MAP files are incompatible");

		reader.close();

		return masks;
	}

	
	public boolean readBlockFile (String blockFile, SNPInfo[] snpList) throws IOException, PlinkFileException
	{		
		
		System.out.println("Searching for BLOCK file: " + blockFile + "...");
		
		try {
			BufferedReader reader = new BufferedReader (new FileReader(blockFile));

			String text = null;

			int blockNo = 0;
			int lastEndIndex = -1;

			while ( (text = reader.readLine()) != null) 
			{			
				String[] blocklist = text.split("\\s+");
				int startIndex = Integer.parseInt( blocklist[0] ) - 1;
				int endIndex =  Integer.parseInt( blocklist[blocklist.length - 1] ) - 1;

				for(int filler=lastEndIndex+1; filler<startIndex; filler++) {
					snpList[filler].setBlockNo(blockNo);
					blockNo ++;
				}

				for(int i=startIndex; i<=endIndex; i++) 
					snpList[i].setBlockNo(blockNo);

				lastEndIndex = endIndex;
				blockNo++;
			}

			reader.close();
		}
		catch (IOException e) {
			System.out.println("Could not find compatible block file. Using default initialization.");
			return false;
		}

		return true;
	}

	
	public void readBedFile(String rawFile) throws IOException
	{
		//TODO
	}

	
	public void readFamFile(String rawFile) throws IOException
	{
		//TODO
	}

}
